from DEAD.AutoDecoder.Evaluation.Interpolation import generate_point_cloud
from DEAD.AutoDecoder.Evaluation.Decoding import get_pytorch_items_from_geometry_vector
from DEAD.Gradiant.AutoGrad import get_vector_field_component as gvc
from DEAD.Gradiant.AutoGrad import auto_grad
from DDPM.Evaluation.InternalBallistic import convert_av1_to_value
from OPT.Config import *
import numpy as np
import torch
import matplotlib.pyplot as plt
import DEAD.AutoDecoder.Evaluation.InternalBallistic as IB
import csv

def objective_function(dead, control_params, t_array_target, p_array_target, cstar, rho, n, At, lv, gv, av1):
    enable_slot_term = control_params["enable_slot_term"]
    slot_min_web = control_params["slot_min_web"]

    enable_port_term = control_params["enable_port_term"]
    central_port_radius = control_params["central_port_radius"]
    port_lz_factor = control_params["port_lz_factor"]

    fator_eta_term = control_params["factor_eta_term"]
    offset = control_params["offset_eta"]

    cstar_rho = cstar*rho

    Ab_array_norm, w_array_norm, w_pred, gradients_norm, V, lz=\
    decodeToApproximateBurningSurface(dead, lv, gv, need_the_gradient_of_gradient=True)

    # =================程函正则项====================
    eikonal_term = 1000*torch.mean((gradients_norm - 1.0) ** 2)

    # =================装填分数计算====================
    loading_fraction, loading_fraction_offset = calculate_loading_fraction(w_pred,0.0)
    eta_term = fator_eta_term*(1.0-loading_fraction_offset) ** 2

    # =================内弹道匹配项================
    r_ref, R = convert_av1_to_value(av1[0])
    t_array_designed, p_array_designed = IB.internal_ballistic_calculation_pytorch(
        w_array_norm=w_array_norm,
        Ab_array_norm=Ab_array_norm,
        cstar_rho=cstar_rho,
        r_ref=r_ref,
        n=n,
        R=R,
        At=At,
        verify_range=False)
    
    r2, r2_k, t_vector_designed, p_vector_designed=\
        performance_matching_degree(t_array_designed,
                                    p_array_designed,
                                    t_array_target,
                                    p_array_target)
    r2_term = (1-r2)*10.0

    # ==================lz条件========================
    lz_term= 100 *torch.clamp(2-lz, min=0.0) ** 2

    # =================喷管安装条件====================
    num_points = 100
    # x, y坐标全为0
    zero_coords = torch.zeros(num_points, device="cuda")

    # =================槽底肉厚条件====================
    # 在0.0到lz之间等距生成z坐标
    z_coords = torch.linspace(0.0, lz,
                              num_points, device="cuda")
    x_coords = torch.ones(num_points, device="cuda")
    # 组合成点云(这个点云沿着z轴分布)
    point_cloud = torch.stack(
        [x_coords, zero_coords, z_coords], dim=1).unsqueeze(0)
    # 预测u值
    u_pred_slot = dead.forward(point_cloud, lv, gv)
    slot_term = 100 *enable_slot_term * torch.mean(torch.clamp((slot_min_web-u_pred_slot)/slot_min_web, min=0.0) ** 2)

    # ==================槽宽条件====================
    if offset>0.0:
        z_coords = torch.linspace(0.0, lz, num_points, device="cuda")
        z_coords = torch.repeat_interleave(z_coords, repeats=num_points)
        x_coords = torch.linspace(0.0, 1.0, num_points, device="cuda")
        x_coords = torch.tile(x_coords, dims=(num_points,))
        zero_coords_expand = torch.zeros(num_points*num_points, device="cuda")
        # 组合成点云(这个点云沿着x_z平面分布)
        point_cloud = torch.stack(
            [x_coords, zero_coords_expand, z_coords], dim=1).unsqueeze(0)
        point_cloud.requires_grad = True
        u_pred_width = offset + dead.forward(point_cloud, lv, gv)
        # 修改部分：当 u_pred_width > offset 时为 0，否则保留原值
        u_pred_width = torch.where(
            u_pred_width > offset,
            torch.zeros_like(u_pred_width),
            u_pred_width
        )
        du_dxyz = auto_grad(u_pred_width, point_cloud, 0, False)
        du_dy=gvc(du_dxyz, 1)
        width_term = 1e3 * torch.mean((torch.clamp(du_dy * u_pred_width / (offset), min=0.0)) ** 2)
    else:
        width_term = torch.tensor(0.0).cuda()

    
    # =================中心流道条件====================
    # 在lz*port_lz_factor到lz之间等距生成z坐标
    z_coords = torch.linspace(lz*port_lz_factor, lz, num_points, device="cuda")
    # 组合成点云(这个点云沿着z轴分布)
    point_cloud = torch.stack(
        [zero_coords, zero_coords, z_coords], dim=1).unsqueeze(0)
    # 预测u值
    u_pred_port = dead.forward(point_cloud, lv, gv)
    port_term = 100 * \
        enable_port_term*torch.mean(torch.clamp((u_pred_port + central_port_radius)/central_port_radius, min=0.0) ** 2)

    # =================计算损失函数（共5项）====================
    objective_value = eikonal_term + eta_term + r2_term + port_term + slot_term + lz_term + width_term
    # 打开一个CSV文件用于追加写入
    with open(f'{control_params["result_path"]}/log.csv', 'a', newline='') as csvfile:
        writer = csv.writer(csvfile)
        
        # 写入表头（仅第一次需要）
        if csvfile.tell() == 0:
            writer.writerow([
                "Objective Value", "R2 Term", "Eta Term", "Eikonal Term", "Slot Term", "Port Term","Width_term", "lz_term", "Loading Fraction"
            ])
        
        # 写入数据行
        writer.writerow([
            objective_value.item(), r2_term.item(), eta_term.item(), eikonal_term.item(), slot_term.item(), port_term.item(),width_term.item(), lz_term.item(), loading_fraction.item()
        ])
    return objective_value, t_vector_designed, p_vector_designed, r2, loading_fraction


def decodeToApproximateBurningSurface(dead, lv, gv, need_the_gradient_of_gradient=True):
    # =================获取几何参数================
    lz, n_slots, m = get_pytorch_items_from_geometry_vector(dead, gv)#注意，lv需要携带梯度信息
    n_slots = np.round(np.clip(n_slots, 2, 15))  # 圆整
    m = np.clip(m, 1.0, 2.0)  # 将m限制在1.0到2.0之间
    V = np.pi*lz-(np.pi*2.0/m-4.0/3.0*np.pi/m)  # 计算空腔与固体区域总体积

    # ========点云生成，预测烧去肉厚场，计算梯度========
    # 生成点云
    point_cloud = generate_point_cloud(
        450000, lz, n_slots, m, max_batch_size=100000).unsqueeze(0)#如果显存不够，可以减小
    point_cloud.requires_grad_(True)  # 计算梯度务必打开
    # 预测w值
    w_pred = dead.forward(point_cloud, lv, gv)
    # 如果要计算关于gradient的梯度，需要打开
    du_dxyz = auto_grad(w_pred, point_cloud, 0, need_the_gradient_of_gradient)
    gradients_norm = torch.sqrt(
        gvc(du_dxyz, 0)**2+gvc(du_dxyz, 1)**2+gvc(du_dxyz, 2)**2)
    # del du_dxyz, point_cloud  # 显式释放显存
    # torch.cuda.empty_cache()  # 清空缓存（可选）
    Ab_array_norm, w_array_norm = compute_isosurface_areas(
        V, w_pred, gradients_norm)
    return Ab_array_norm, w_array_norm, w_pred, gradients_norm, V, lz
def compute_isosurface_areas(V, values, gradients_norm, num_isosurfaces=50):
    """
    计算N个等值面的面积（隐式曲面 + 蒙特卡洛积分）
    适配输入形状： values: (1, N, 1)

    参数：
        points:  (1, N, 3) 点云坐标
        values:  (1, N, 1) 点云对应的u值
        gradients_norm: (1, N, 1) 点云对应的梯度模长
        num_isosurfaces: 等值面数量（默认50）

    返回：
        areas: (num_isosurfaces,) 每个等值面的面积
        c_values: (num_isosurfaces,) 等值面对应的u值
    """
    device = values.device

    # 1. 移除批次维度并确保形状正确
    values = values.squeeze(0).squeeze(-1)  # (N,)
    gradients_norm = gradients_norm.squeeze(0).squeeze(-1)  # (N,)

    # 2. 生成等值面的目标u值（均匀分布在u的[min, max]范围内）
    u_min = 0
    u_max = torch.max(values).item()*0.999
    c_values = torch.linspace(
        u_min, u_max, num_isosurfaces, device=device)  # (num_isosurfaces,)
    # sigma = (u_max - u_min)*0.03
    sigma = (u_max - u_min) / (num_isosurfaces * 1.0)  # 调整系数以优化效果

    # 3. 高斯核近似 Dirac delta 函数，并乘以梯度模长
    residuals = values.unsqueeze(
        1) - c_values.unsqueeze(0)  # (N, num_isosurfaces)
    delta = torch.exp(-0.5 * (residuals / sigma)**2) / \
        (sigma * np.sqrt(2 * np.pi))
    delta = delta * gradients_norm.unsqueeze(1)      # (N, num_isosurfaces)

    # 4. 蒙特卡洛积分（假设梯度模长=1）
    areas = V * torch.mean(delta, dim=0)  # (num_isosurfaces,)
    areas[-1] = torch.tensor([0.0])  # 注意最后一个点通常有较大误差，因此将其置为0

    return areas, c_values

def calculate_loading_fraction(w_pred, offset):
    # 计算正w值的比例(体积装填分数)
    loading_fraction_offset = torch.mean(
        torch.sigmoid((w_pred+offset) * 1e4))  # offset是用于防止优化后得到狭小的缝
    with torch.no_grad():
        loading_fraction = torch.mean(torch.sigmoid((w_pred) * 1e9))
    return loading_fraction, loading_fraction_offset

def performance_matching_degree(t_array_designed, p_array_designed, t_array_target,p_array_target,resolution=512):
    # 计算性能匹配度
    # 计算共同时间上限（乘以1.2以确保计算的稳定性）
    t_max = torch.max(t_array_target[-1], t_array_designed[-1])*1.2
    _, t_vector_designed, p_vector_designed = IB.vectoring_internal_ballistic_pytorch(
        t_array_designed, p_array_designed, t_max, resolution)
    _, _, p_vector_target = IB.vectoring_internal_ballistic_pytorch(
        t_array_target, p_array_target, t_max, resolution)
    # 计算r2以评估内弹道目标之间的差异
    r2, r2_k = calculate_r2(p_vector_designed, p_vector_target)
    return r2, r2_k, t_vector_designed, p_vector_designed

def calculate_r2(p_vector_designed, p_vector_target):
    # 计算目标值的均值
    y_mean = torch.mean(p_vector_target)
    k=torch.sum(p_vector_target*p_vector_designed)/torch.sum(p_vector_designed*p_vector_designed)
    # 计算残差平方和 (SSR)
    ssr = torch.sum((p_vector_target - p_vector_designed) ** 2)
    ssr_k=torch.sum((p_vector_target - k*p_vector_designed) ** 2)
    # 计算总平方和 (SST)
    sst = torch.sum((p_vector_target - y_mean) ** 2)
    # 计算 R²
    r2 = 1 - (ssr / sst)
    r2_k = 1 - (ssr_k / sst)
    return r2, r2_k  # 返回标量值

def plot_isosurface_areas(c_values, areas):
    """
    绘制等值面面积随 u 值变化的曲线
    参数：
        c_values: (num_isosurfaces,) 等值面的 u 值
        areas: (num_isosurfaces,) 对应的等值面面积
    """
    plt.figure(figsize=(10, 6))
    plt.plot(c_values.cpu().numpy(), areas.cpu().numpy(),
             'b-', linewidth=2, label="Area vs. Isovalue")
    plt.xlabel("Isovalue (c)", fontsize=12)
    plt.ylabel("Area", fontsize=12)
    plt.title("Isosurface Area as a Function of Isovalue", fontsize=14)
    plt.grid(True, linestyle='--', alpha=0.7)
    plt.legend()
    plt.show()


def plot_isosurface_areas_compare(c_values, areas, std_w, std_Ab):
    """
    绘制等值面面积随 u 值变化的曲线
    参数：
        c_values: (num_isosurfaces,) 等值面的 u 值
        areas: (num_isosurfaces,) 对应的等值面面积
    """
    plt.figure(figsize=(10, 6))
    plt.plot(c_values.cpu().numpy(), areas.cpu().numpy(),
             'b-', linewidth=2, label="Random method")
    plt.plot(std_w, std_Ab, 'r-', linewidth=2, label="VTK method")
    plt.xlabel("Isovalue (c)", fontsize=12)
    plt.ylabel("Area", fontsize=12)
    plt.title("Isosurface Area as a Function of Isovalue", fontsize=14)
    plt.grid(True, linestyle='--', alpha=0.7)
    plt.legend()
    plt.show()